"""
target detection
"""
import rospy
import cv2
from std_msgs.msg import String
from sensor_msgs.msg import Image
from swarm_msgs.msg import BoundingBox, BoundingBoxes
import os
import re
import numpy as np
import platform
import time
from datetime import datetime
import threading
from datetime import datetime
import subprocess

rospy.init_node('my_node_name', anonymous=True)
from YoloRknn import YoloRknn

yolo = YoloRknn('/home/orangepi/zny_ws/src/wlcb_exploration/relative_pkg/det/scripts/ysl.rknn')
count = 0

def clear_folder(folder_path):
    """
    删除 folder_path 下的所有文件和子目录，但保留 folder_path 本身。
    """
    for entry in os.listdir(folder_path):
        entry_path = os.path.join(folder_path, entry)
        # 如果是文件或链接，直接删除
        if os.path.isfile(entry_path) or os.path.islink(entry_path):
            os.remove(entry_path)
        # 如果是目录，递归删除整个目录
        elif os.path.isdir(entry_path):
            shutil.rmtree(entry_path)

class FrameReader(threading.Thread):
    def __init__(self):
        threading.Thread.__init__(self)
        self.cap = cv2.VideoCapture("/dev/down_camera")  # 从摄像头读取视频
        self.frame = None
        self.running = True
        self.lock = threading.Lock()
        self.human_bbox_pub = rospy.Publisher("/bian/human", BoundingBoxes, queue_size=10)
        self.car_bbox_pub = rospy.Publisher("/bian/car", BoundingBoxes, queue_size=10)
        self.oil_bbox_pub = rospy.Publisher("/bian/oil", BoundingBoxes, queue_size=10)
        self.weapon_bbox_pub = rospy.Publisher("/bian/weapon", BoundingBoxes, queue_size=10)

        self.bigarmor_bbox_pub = rospy.Publisher("/bian/bigarmor", BoundingBoxes, queue_size=10)
        self.shatanche_bbox_pub = rospy.Publisher("/bian/shatanche", BoundingBoxes, queue_size=10)
        self.t_pub = rospy.Publisher("/img", BoundingBoxes, queue_size=10)
        
        if not self.cap.isOpened():
            raise ValueError("Could not open video device")

    def run(self):
        while self.running:
            self.cap.grab()
            ret, frame = self.cap.retrieve()
            with self.lock:
                if ret:
                    self.frame = frame.copy()

                else:
                    self.frame = None
            time.sleep(0.01)

    def stop(self):
        self.running = False
        self.cap.release()


class YoloProcessor(threading.Thread):
    def __init__(self, frame_reader):
        threading.Thread.__init__(self)
        self.frame_reader = frame_reader
        self.write_count = 0
        self.running = True
        self.bboxes = BoundingBoxes()
        self.human_bboxes = BoundingBoxes()
        self.car_bboxes = BoundingBoxes()
        self.oil_bboxes = BoundingBoxes()
        self.weapon_bboxes = BoundingBoxes()
        self.bigarmor_bboxes = BoundingBoxes()
        self.shatanche_bboxes = BoundingBoxes()
        self.org_img = np.zeros((552, 640, 3), dtype=np.uint8)  # 创建一个空图像

        self.save_count = 0
        self.save_path = '/home/orangepi/img'
        os.makedirs(self.save_path, exist_ok=True)
        clear_folder(self.save_path)

    def run(self):
        while self.running:
            frame = None
            with self.frame_reader.lock:
                if self.frame_reader.frame is not None:
                    frame = self.frame_reader.frame.copy()
                    self.frame_reader.t_pub.publish(self.bboxes)
            if frame is not None:
                self.process_frame(frame)
            #time.sleep(0.01)  # 控制处理帧的速度

    def process_frame(self, frame):
        try:
            self.human_bboxes.bounding_boxes = []
            self.car_bboxes.bounding_boxes = []
            self.oil_bboxes.bounding_boxes = []
            self.weapon_bboxes.bounding_boxes = []
            self.bigarmor_bboxes.bounding_boxes = []
            self.shatanche_bboxes.bounding_boxes = []


            yolo.det(frame)

            
            # img_test = cv2.resize(frame, (1024, 768))
            # cv2.imshow("test", img_test)
            # cv2.waitKey(1)

            
            current_time = datetime.now().strftime("%Y%m%d_%H%M%S")     
                # 构建文件名
            filename = f"{self.save_count}_{current_time}.jpg"
                # 保存图片
            # os.makedirs('img', exist_ok=True)
            save_full_path = os.path.join(self.save_path, filename)
            if len(yolo.det_bboxes) != 0:  # 如果检测到了目标
            # 获取当前时间并格式化为字符串
                cv2.imwrite(save_full_path, yolo.res_img)
                # cv2.imwrite(save_full_path, yolo.res_img)
                # 计数器加1
            else :
                cv2.imwrite(save_full_path, yolo.org_img)
            self.save_count += 1
        
            print(f"已保存图片: {filename}")
            
            for det_bbox in yolo.det_bboxes:
                bbox = BoundingBox()
                bbox.xmin = det_bbox['left']
                bbox.ymin = det_bbox['top']
                bbox.xmax = det_bbox['right']
                bbox.ymax = det_bbox['bottom']
                bbox.Class = det_bbox['id']
                if det_bbox['score']<0.45:
                    continue
                # if det_bbox['score'] > 0.7:
                #     self.save_count = self.save_count + 1
                #     save_name = f"{det_bbox['id']}_{self.save_count}_{det_bbox['score']:.2f}.jpg"
                #     save_full_path = os.path.join(self.save_path, save_name)
                #     try:
                #         cv2.imwrite(save_full_path, yolo.res_img)
                #         print(f"Saved detection result: {save_full_path}")
                #         # self.save_count[det_bbox['id']] = self.save_count.get(det_bbox['id'], 0) + 1
                #     except Exception as e:
                #         print(f"Failed to save image: {e}")

                # print(bbox.Class)
                if bbox.Class == 'person':
                    self.human_bboxes.bounding_boxes.append(bbox)
                    # print(bbox)
                if bbox.Class == 'car':
                    self.car_bboxes.bounding_boxes.append(bbox)
                    # print(bbox)
                if bbox.Class == 'oil':
                    self.oil_bboxes.bounding_boxes.append(bbox)
                    # print(bbox)
                if bbox.Class == 'armor':
                    self.weapon_bboxes.bounding_boxes.append(bbox)
                    # print(bbox)
                if bbox.Class == 'bigarmor':
                    self.bigarmor_bboxes.bounding_boxes.append(bbox)
                    # print(bbox)
                if bbox.Class == 'shatanche':
                    self.shatanche_bboxes.bounding_boxes.append(bbox)
                    # print(bbox)
                yolo.det_bboxes = []
            
            if len(self.human_bboxes.bounding_boxes) > 0:
                try:
                    self.frame_reader.human_bbox_pub.publish(self.human_bboxes)
                except Exception as e:
                    print(f"An 1error occurred: {e}")
            if len(self.car_bboxes.bounding_boxes) > 0:
                try:
                    self.frame_reader.car_bbox_pub.publish(self.car_bboxes)
                except Exception as e:
                    print(f"An 1error occurred: {e}")
            if len(self.oil_bboxes.bounding_boxes) > 0:
                try:
                    self.frame_reader.oil_bbox_pub.publish(self.oil_bboxes)
                except Exception as e:
                    print(f"An 1error occurred: {e}")
            if len(self.weapon_bboxes.bounding_boxes) > 0:
                try:
                    self.frame_reader.weapon_bbox_pub.publish(self.weapon_bboxes)
                except Exception as e:
                    print(f"An 1error occurred: {e}")       
            if len(self.bigarmor_bboxes.bounding_boxes) > 0:
                try:
                    self.frame_reader.bigarmor_bbox_pub.publish(self.bigarmor_bboxes)
                except Exception as e:
                    print(f"An 1error occurred: {e}")  
            if len(self.shatanche_bboxes.bounding_boxes) > 0:
                try:
                    self.frame_reader.shatanche_bbox_pub.publish(self.shatanche_bboxes)
                except Exception as e:
                    print(f"An 1error occurred: {e}")  
            
            img_1 = cv2.resize(yolo.res_img, (640, 552))
            # if img_1.size > 0:
            #     #创建窗口
            #     cv2.namedWindow('horizonal_cam window')
            #     # 移动窗口到(100, 100)
            #     cv2.moveWindow('horizonal_cam window', 100, 100)
            #     cv2.imshow("horizonal_cam window", img_1)
            #     cv2.waitKey(1)
                
            #     if self.write_count%5==0:
                    
            #         save_dir = '/home/orangepi/pictures_yolo'
            #         if not os.path.exists(save_dir):
            #             os.makedirs(save_dir)
            #         now = datetime.now()
            #         timestamp = now.strftime('%Y%m%d_%H%M%S')
            #         filename = os.path.join(save_dir, f'{timestamp}.png')
            #         cv2.imwrite(filename, img_1)
            
            # self.write_count+=1
        except Exception as e:
            print(f"An 2error occurred: {e}")

    def stop(self):
        self.running = False

def main():

    frame_reader = FrameReader()
    yolo_processor = YoloProcessor(frame_reader)

    frame_reader.start()
    yolo_processor.start()

   
    try:
        print("1")
        # while True:
        #     frame = None
        #     with frame_reader.lock:
        #         if frame_reader.frame is not None:
        #             frame = frame_reader.frame.copy()
        #
        #     if frame is not None:
        #         cv2.imshow('Frame', frame)
        #         if cv2.waitKey(1) & 0xFF == ord('q'):
        #             break
        #     time.sleep(0.01)
    except KeyboardInterrupt:
        frame_reader.stop()
        yolo_processor.stop()
    except Exception as e:
        print(f"An error occurred in main: {e}")

    frame_reader.join()
    yolo_processor.join()
    cv2.destroyAllWindows()


if __name__ == '__main__':
    main()
    rospy.spin()
